﻿using ModelHasher.Framework;
using ModelHasher.Utils;
using StableDiffusionTools.Framework;
using StableDiffusionTools.Utils;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ModelHasher
{
    public class ModelHasherState
    {
        private static readonly object locker = new object();

        private static ModelHasherState state;
        public static ModelHasherState State
        {
            get
            {
                lock (locker)
                {
                    if (state == null)
                        state = new ModelHasherState();
                }

                return state;
            }
        }

        private SettingsManager<ModelHasherSettings> settingsMan = new SettingsManager<ModelHasherSettings>(nameof(ModelHasherSettings), createBackups: true);
        private ModelHasherSettings settings;
        private Dictionary<string, ModelHash> hashToModelLegacy = new Dictionary<string, ModelHash>();
        private Dictionary<string, ModelHash> hashToModelShort = new Dictionary<string, ModelHash>();
        private Dictionary<string, ModelHash> hashToModelSHA256 = new Dictionary<string, ModelHash>();
        private object stateLocker = new object();

        private ModelHasherState()
        {
            settings = settingsMan.LoadSettings();
            for (int i = 0; i < settings.Models.Count; i++)
            {
                AddToDicts(settings.Models[i]);
            }
        }

        public bool TryGetModelHashByPathAndName(string name, out ModelHash modelHash)
        {
            lock (stateLocker)
            {
                modelHash = settings.Models.Find(m => m.Name.Equals(name));
            }
            return modelHash != null;
        }

        public bool TryGetModelHashByName(string name, out ModelHash modelHash)
        {
            lock (stateLocker)
            {
                modelHash = settings.Models.Find(m => ModelHashPathAndNameMatchesName(m.Name, name));
            }
            return modelHash != null;
        }

        public bool TryGetModelHashByHash(string hash, out ModelHash modelHash)
        {
            lock (stateLocker)
            {
                if (hash.Length == 8)
                    hashToModelLegacy.TryGetValue(hash, out modelHash);
                else if (hash.Length == 10)
                    hashToModelShort.TryGetValue(hash, out modelHash);
                else
                    hashToModelSHA256.TryGetValue(hash, out modelHash);

            }
            return modelHash != null;
        }

        public async Task<ModelHash> UpdateStateFromPathAsync(string path, string root, Progress<HashingUtil.HasherProgress> progress, bool forced)
        {
            ModelHash result = null;
            DirectoryInfo rootDir = new DirectoryInfo(root);
            PathInfo info = ExtendedPathInfo.GetPathInfo(path);
            if(info == PathInfo.File)
            {
                FileInfo fileInfo = new FileInfo(path);
                string matchName = fileInfo.FullName.Substring(rootDir.FullName.Length + 1);
                bool foundMatch = TryGetModelHashByPathAndName(matchName, out ModelHash existingModelHash);
                if (foundMatch)
                {
                    result = existingModelHash;
                }

                if (!foundMatch)
                {
                    string hash = null;
                    try
                    {
                        hash = await HashingUtil.CalculateFileHashAsync(HashAlgo.SHA256, path, progress: progress, lowerCase: true);
                    }
                    catch (Exception) { }

                    result = new ModelHash(matchName, hash);

                    lock (stateLocker)
                    {
                        settings.Models.Add(result);
                        AddToDicts(result);
                    }
                }
                else if (forced)
                {
                    string hash = null;
                    try
                    {
                        hash = await HashingUtil.CalculateFileHashAsync(HashAlgo.SHA256, path, progress: progress, lowerCase: true);

                        lock (stateLocker)
                        {
                            var shortHash = hash.Substring(0, 10);
                            if (!result.ShortHash.Equals(shortHash))
                            {
                                RemoveFromDicts(result);
                                result.UpdateHash(hash);
                                AddToDicts(result);
                            }
                        }
                    }
                    catch (Exception) { }
                }
            }
            return result;
        }

        public void UpdateModelHashName(ModelHash modelHash, string newName)
        {
            RemoveFromDicts(modelHash);
            modelHash.Name = newName;
            AddToDicts(modelHash);
        }

        public void SaveSettings()
        {
            lock (stateLocker)
            {
                settingsMan.SaveSettings(settings);
            }
        }

        private void AddToDicts(ModelHash model)
        {
            hashToModelLegacy[model.LegacyHash] = model;
            hashToModelShort[model.ShortHash] = model;
            hashToModelSHA256[model.Hash] = model;
        }

        private void RemoveFromDicts(ModelHash model)
        {
            hashToModelLegacy.Remove(model.LegacyHash);
            hashToModelShort.Remove(model.ShortHash);
            hashToModelSHA256.Remove(model.Hash);
        }

        public static bool ModelExistsOnDisk(string fullPath)
        {
            bool exists = false;
            try
            {
                exists = ExtendedPathInfo.GetPathInfo(fullPath) == PathInfo.File;
            }
            catch (Exception) { }
            return exists;
        }

        private static bool ModelHashPathAndNameMatchesName(string storedModelHashPathAndName, string matchName)
        {
            int lastSlashIndex = storedModelHashPathAndName.LastIndexOf('\\');
            if (lastSlashIndex != -1)
            {
                string storedName = storedModelHashPathAndName.Substring(lastSlashIndex + 1);
                return storedName.Equals(matchName);
            }
            return false;
        }
    }
}
